import pickle
import random

import matplotlib.pyplot as plt
import numpy.random
from sklearn.manifold import TSNE

output_file = "affectnet8class.pkl"
label_file = "affectnet8class_lable.pkl"

with open(output_file, 'rb') as f:
    output = pickle.load(f)

with open(label_file, 'rb') as f:
    labels = pickle.load(f)

size = list(range(len(output)))
select = random.sample(size, 372)
output = [output[i] for i in select]
labels = [labels[i] for i in select]

index = ["SU", "FE", "DI", "HA", "SA", "AN", "NE", "CO"]
new_o = []
new_l = []
for i in range(len(labels)):
    if labels[i] == 7:
        continue
    new_o.append(output[i])
    new_l.append(labels[i])

tsne = TSNE(n_components=2, init='pca', random_state=numpy.random, perplexity=50, learning_rate=100)
embedded_features = tsne.fit_transform(new_o)

# 绘制tSNE图，并根据类别标签使用不同的颜色区分数据点
plt.figure(figsize=(8, 6))
scatter = plt.scatter(embedded_features[:, 0], embedded_features[:, 1],
                      s=20,
                      c=new_l,
                      cmap='viridis',
                      )
plt.colorbar(scatter, ticks=range(7))
plt.title('tSNE Plot of Emotion Features')
plt.legend(handles=scatter.legend_elements()[0], labels=index)
plt.show()
